use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use crate::context::Context;
use crate::errors::SendError;
use crate::message::DynamicMessage;
use crate::middleware::Middleware;

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum CircuitBreakerState {
    Closed,
    HalfOpen,
    Open,
}

#[derive(Debug, Clone)]
pub struct CircuitBreakerConfig {
    pub failure_threshold: u32,
    pub reset_timeout: Duration,
    pub half_open_timeout: Duration,
}

impl Default for CircuitBreakerConfig {
    fn default() -> Self {
        Self {
            failure_threshold: 5,
            reset_timeout: Duration::from_secs(60),
            half_open_timeout: Duration::from_secs(30),
        }
    }
}

pub struct CircuitBreaker {
    state: Arc<RwLock<CircuitBreakerState>>,
    failures: Arc<RwLock<u32>>,
    last_failure: Arc<RwLock<Option<Instant>>>,
    config: CircuitBreakerConfig,
}

impl CircuitBreaker {
    pub fn new(config: CircuitBreakerConfig) -> Self {
        Self {
            state: Arc::new(RwLock::new(CircuitBreakerState::Closed)),
            failures: Arc::new(RwLock::new(0)),
            last_failure: Arc::new(RwLock::new(None)),
            config,
        }
    }

    async fn record_failure(&self) {
        let mut failures = self.failures.write().await;
        *failures += 1;
        *self.last_failure.write().await = Some(Instant::now());

        if *failures >= self.config.failure_threshold {
            *self.state.write().await = CircuitBreakerState::Open;
        }
    }

    async fn record_success(&self) {
        let mut failures = self.failures.write().await;
        *failures = 0;
        *self.state.write().await = CircuitBreakerState::Closed;
    }

    async fn check_state(&self) -> CircuitBreakerState {
        let state = *self.state.read().await;
        match state {
            CircuitBreakerState::Open => {
                if let Some(last_failure) = *self.last_failure.read().await {
                    if last_failure.elapsed() >= self.config.reset_timeout {
                        *self.state.write().await = CircuitBreakerState::HalfOpen;
                        return CircuitBreakerState::HalfOpen;
                    }
                }
            }
            CircuitBreakerState::HalfOpen => {
                if let Some(last_failure) = *self.last_failure.read().await {
                    if last_failure.elapsed() >= self.config.half_open_timeout {
                        *self.state.write().await = CircuitBreakerState::Closed;
                        return CircuitBreakerState::Closed;
                    }
                }
            }
            CircuitBreakerState::Closed => {}
        }
        state
    }
}

impl Middleware for CircuitBreaker {
    fn handle<'a>(&'a self, ctx: &'a mut Context, msg: DynamicMessage, next: Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = Result<(), SendError>> + Send + 'a>> + Send + 'a>) -> Pin<Box<dyn Future<Output = Result<(), SendError>> + Send + 'a>> {
        Box::pin(async move {
            match self.check_state().await {
                CircuitBreakerState::Open => {
                    Err(SendError::CircuitBreakerOpen)
                }
                CircuitBreakerState::HalfOpen | CircuitBreakerState::Closed => {
                    match next().await {
                        Ok(_) => {
                            self.record_success().await;
                            Ok(())
                        }
                        Err(e) => {
                            self.record_failure().await;
                            Err(e)
                        }
                    }
                }
            }
        })
    }
} 